/*
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.firebase.ai.type

import kotlinx.serialization.DeserializationStrategy
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonContentPolymorphicSerializer
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonNull
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.jsonObject

/**
 * Parent interface for responses from the model during live interactions.
 *
 * @see LiveServerContent
 * @see LiveServerToolCall
 * @see LiveServerToolCallCancellation
 * @see LiveServerSetupComplete
 */
@PublicPreviewAPI public interface LiveServerMessage

/**
 * Incremental server update generated by the model in response to client messages.
 *
 * Content is generated as quickly as possible, and not in realtime. You may choose to buffer and
 * play it out in realtime.
 */
@PublicPreviewAPI
public class LiveServerContent
@Deprecated("This class should not be constructed, only received from the Server")
public constructor(
  /**
   * The content that the model has generated as part of the current conversation with the user.
   *
   * This can be `null` if there is no content.
   */
  public val content: Content?,

  /**
   * The model was interrupted by the client while generating data.
   *
   * An interruption occurs when the client sends a message while the model is actively sending
   * data.
   */
  public val interrupted: Boolean,

  /**
   * The model has finished sending data in the current turn.
   *
   * Generation will only start in response to additional client messages.
   *
   * Can be set alongside [content], indicating that the [content] is the last in the turn.
   *
   * @see generationComplete
   */
  public val turnComplete: Boolean,

  /**
   * The model has finished _generating_ data for the current turn.
   *
   * For realtime playback, there will be a delay between when the model finishes generating content
   * and the client has finished playing back the generated content. [generationComplete] indicates
   * that the model is done generating data, while [turnComplete] indicates the model is waiting for
   * additional client messages. Sending a message during this delay may cause an [interrupted]
   * message to be sent.
   *
   * Note that if the model was [interrupted], this will not be set. The model will go from
   * [interrupted] -> [turnComplete].
   */
  public val generationComplete: Boolean,

  /**
   * The input transcription. The transcription is independent to the model turn which means it
   * doesn't imply any ordering between transcription and model turn.
   */
  public val inputTranscription: Transcription?,

  /**
   * The output transcription. The transcription is independent to the model turn which means it
   * doesn't imply any ordering between transcription and model turn.
   */
  public val outputTranscription: Transcription?
) : LiveServerMessage {
  @OptIn(ExperimentalSerializationApi::class)
  @Serializable
  internal data class Internal(
    val modelTurn: Content.Internal?,
    val interrupted: Boolean?,
    val turnComplete: Boolean?,
    val generationComplete: Boolean?,
    val inputTranscription: Transcription.Internal?,
    val outputTranscription: Transcription.Internal?
  )
  @Serializable
  internal data class InternalWrapper(val serverContent: Internal) : InternalLiveServerMessage {
    @OptIn(ExperimentalSerializationApi::class)
    override fun toPublic(): LiveServerContent {
      // WhenMajor(Revisit the decision to make these have default values)
      return LiveServerContent(
        serverContent.modelTurn?.toPublic(),
        serverContent.interrupted ?: false,
        serverContent.turnComplete ?: false,
        serverContent.generationComplete ?: false,
        serverContent.inputTranscription?.toPublic(),
        serverContent.outputTranscription?.toPublic()
      )
    }
  }
}

/** The model is ready to receive client messages. */
@PublicPreviewAPI
public class LiveServerSetupComplete : LiveServerMessage {
  @Serializable
  internal data class Internal(val setupComplete: JsonObject) : InternalLiveServerMessage {
    override fun toPublic() = LiveServerSetupComplete()
  }
}

/**
 * Request for the client to execute the provided [functionCalls].
 *
 * The client should return matching [FunctionResponsePart], where the `id` fields correspond to
 * individual [FunctionCallPart]s.
 *
 * @property functionCalls A list of [FunctionCallPart] to run and return responses for.
 */
@PublicPreviewAPI
public class LiveServerToolCall(public val functionCalls: List<FunctionCallPart>) :
  LiveServerMessage {
  @Serializable
  internal data class Internal(
    val functionCalls: List<FunctionCallPart.Internal.FunctionCall> = emptyList()
  )
  @Serializable
  internal data class InternalWrapper(val toolCall: Internal) : InternalLiveServerMessage {
    override fun toPublic() =
      LiveServerToolCall(
        toolCall.functionCalls.map { functionCall ->
          FunctionCallPart(
            name = functionCall.name,
            args = functionCall.args.orEmpty().mapValues { it.value ?: JsonNull },
            id = functionCall.id
          )
        }
      )
  }
}

/**
 * Notification for the client to cancel a previous function call from [LiveServerToolCall].
 *
 * You do not need to send [FunctionResponsePart]s for the cancelled [FunctionCallPart]s.
 *
 * @property functionIds A list of `id`s matching the `id` provided in a previous
 * [LiveServerToolCall], where only the provided `id`s should be cancelled.
 */
@PublicPreviewAPI
public class LiveServerToolCallCancellation(public val functionIds: List<String>) :
  LiveServerMessage {
  @Serializable internal data class Internal(val functionIds: List<String> = emptyList())
  @Serializable
  internal data class InternalWrapper(val toolCallCancellation: Internal) :
    InternalLiveServerMessage {
    override fun toPublic() = LiveServerToolCallCancellation(toolCallCancellation.functionIds)
  }
}

@PublicPreviewAPI
@Serializable(LiveServerMessageSerializer::class)
internal sealed interface InternalLiveServerMessage {
  fun toPublic(): LiveServerMessage
}

@OptIn(PublicPreviewAPI::class)
internal object LiveServerMessageSerializer :
  JsonContentPolymorphicSerializer<InternalLiveServerMessage>(InternalLiveServerMessage::class) {
  @OptIn(PublicPreviewAPI::class)
  override fun selectDeserializer(
    element: JsonElement
  ): DeserializationStrategy<InternalLiveServerMessage> {
    val jsonObject = element.jsonObject
    return when {
      "serverContent" in jsonObject -> LiveServerContent.InternalWrapper.serializer()
      "setupComplete" in jsonObject -> LiveServerSetupComplete.Internal.serializer()
      "toolCall" in jsonObject -> LiveServerToolCall.InternalWrapper.serializer()
      "toolCallCancellation" in jsonObject ->
        LiveServerToolCallCancellation.InternalWrapper.serializer()
      else ->
        throw SerializationException(
          "Unknown LiveServerMessage response type. Keys found: ${jsonObject.keys}"
        )
    }
  }
}
